# $Id: udp.py 9227e81e919a 2012/08/27 15:49:57 pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

__version__  = "$Revision: 9227e81e919a $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2012-08-27 15:49:57Z $"

# essential imports
import socket
import select

# Error class
class SocketError(Exception):
   pass

# Main udp object
class udp(object):
   def __init__(self, Host = None, Port = 0):
      # instantiate the error class
      self._err = SocketError()
      self._sock = None
      self._host = None
      self._port = 0

      if Host:
         self._host = Host

      if Port:
         # check to see whether the port is integer or raise error
         try:
            self._port = int(Port)
         except ValueError, message:
            self._error(message)

      sockList = []
      # Use getaddrinfo to create socket, in this way it will work for
      # both IPv4 and IPv6 families
      try:
         sockList = socket.getaddrinfo(self._host, self._port,
               socket.AF_UNSPEC, socket.SOCK_DGRAM)
      except socket.gaierror, (error, message):
         self._error("getaddrinfo: %s" %message)

      # iterate through the socket list and create the socket fd for
      # the same
      for res in sockList:
         af, socktype, proto, canonname, sa = res
         try:
            self._sock = socket.socket(af, socktype, proto)
         except socket.error, (error, message):
            self._sock = None
            continue

         break

      # If we still have None type for socket fd, then there is an 
      # actual problem, raise error
      if self._sock is None:
         self._error("cannot get a udp socket to %s on port %d" %(
            self._host, self._port))

      if self._host:
         self._addr = (self._host, self._port)
      else:
         # If we are server, then host is just empty string, host cannot be
         # of None type
         self._addr = ("", self._port)

   # Set/Unset blocking based on the flag
   def Blocking(self, Flag):
      if self._sock:
         self._sock.setblocking(Flag)

   # Simple function, listens, i.e. binds to the socket
   def Listen(self):
      if self._sock:
         self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         try:
            self._sock.bind(self._addr)
         except socket.error, (error, message):
            self._error(message)

   # write function, obvious writes data to the socket, when timeout is
   # specified, waits till specified time, otherwise, waits forever, if it is
   # a blocking socket
   def Write(self, Data, Timeout = 0):
      sent = None
      if Data:
         try:
            if Timeout:
               rfd, wfd, efd = select.select([], [self._sock], [], Timeout)
            else:
               rfd, wfd, efd = select.select([], [self._sock], [])
         except select.error, message:
            self._error("select error: %s" %message)

         if self._sock in wfd:
            try:
               sent = self._sock.sendto(Data, self._addr)
            except socket.error, (error, message):
               self._error("socket error: %s" %message)
      else:
         self._error("no data provided to send!")

      # always assert what you sent
      assert(sent == len(Data))
      return sent

   # Same thing as write
   def Read(self, bufSize = 1024, Timeout = 0):
      if not bufSize:
         bufSize = 1024

      data = None
      try:
         if Timeout:
            rfd, wfd, efd = select.select([self._sock], [], [], Timeout)
         else:
            rfd, wfd, efd = select.select([self._sock], [], [])
      except select.error, message:
         self._error("select error: %s" %message)

      if self._sock in rfd:
         try:
            data, self._addr = self._sock.recvfrom(bufSize)
         except socket.error, (error, message):
            self._error("socket error: %s" %message)

      return (data, self._addr)

   def Close(self):
      if self._sock:
         # close the socket handle, called by error function
         # before we error out!
         self._sock.close()
         self._sock = None

   def _error(self, Message):
      self.Close()
      raise SocketError, Message
